package org.lemon.transport.rule.engine.mqtt;

import com.google.common.collect.HashMultimap;
import com.google.common.collect.ImmutableSet;
import io.netty.bootstrap.Bootstrap;
import io.netty.buffer.ByteBuf;
import io.netty.channel.*;
import io.netty.channel.nio.NioEventLoopGroup;
import io.netty.channel.socket.SocketChannel;
import io.netty.handler.codec.mqtt.*;
import io.netty.handler.ssl.SslContext;
import io.netty.handler.timeout.IdleStateHandler;
import io.netty.util.concurrent.DefaultPromise;
import io.netty.util.concurrent.Future;
import io.netty.util.concurrent.Promise;
import lombok.extern.slf4j.Slf4j;
import org.lemon.transport.rule.engine.mqtt.config.MqttClientConfig;
import org.lemon.transport.rule.engine.mqtt.handler.MqttChannelHandler;
import org.lemon.transport.rule.engine.mqtt.handler.MqttClientCallback;
import org.lemon.transport.rule.engine.mqtt.handler.MqttHandler;
import org.lemon.transport.rule.engine.mqtt.handler.MqttPingHandler;
import org.lemon.transport.rule.engine.mqtt.model.*;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

@Slf4j
public class MqttClientImpl implements MqttClient {

    /**
     * 服务订阅的主题
     */
    private final Set<String> serverSubscriptions = new HashSet<>();
    /**
     * ping订阅的主题
     */
    private final Set<String> pendingSubscribeTopics = new HashSet<>();

    private final ConcurrentMap<Integer, MqttPendingSubscription> pendingSubscriptions = new ConcurrentHashMap<>();
    private final ConcurrentMap<Integer, MqttPendingUnsubscription> pendingServerUnsubscribes = new ConcurrentHashMap<>();
    private final ConcurrentMap<Integer, MqttIncomingQos2Publish> qos2PendingIncomingPublishes = new ConcurrentHashMap<>();
    private final ConcurrentMap<Integer, MqttPendingPublish> pendingPublishes = new ConcurrentHashMap<>();

    private final HashMultimap<String, MqttSubscription> subscriptions = HashMultimap.create();
    private final HashMultimap<MqttHandler, MqttSubscription> handlerToSubscription = HashMultimap.create();

    private final AtomicInteger nextMessageId = new AtomicInteger(1);

    private final MqttClientConfig mqttClientConfig;

    private final MqttHandler defaultHandler;

    private EventLoopGroup eventLoop;

    private volatile Channel channel;

    private volatile boolean disconnected = false;
    private volatile boolean reconnect = false;

    private String host;
    private int port;
    private MqttClientCallback callback;

    public MqttClientImpl(MqttHandler defaultHandler) {
        this.mqttClientConfig = new MqttClientConfig();
        this.defaultHandler = defaultHandler;
    }

    public MqttClientImpl(MqttClientConfig config, MqttHandler defaultHandler) {
        this.mqttClientConfig = config;
        this.defaultHandler = defaultHandler;
    }

    @Override
    public Future<MqttConnectResult> connect(String host) {
        return this.connect(host, 1883);
    }

    @Override
    public Future<MqttConnectResult> connect(String host, int port) {
        return this.connect(host, port, false);
    }

    @Override
    public Future<MqttConnectResult> connect(String host, int port, boolean reconnect) {

        log.trace("[{}] Connecting to server, isReconnect - {}", channel != null ? channel.id() : "UNKNOWN", reconnect);
        if (this.eventLoop == null) {
            this.eventLoop = new NioEventLoopGroup();
        }
        Promise<MqttConnectResult> connectFuture = new DefaultPromise<>(this.eventLoop.next());
        Bootstrap bootstrap = new Bootstrap();
        bootstrap.group(this.eventLoop);
        bootstrap.remoteAddress(host, port);
        bootstrap.handler(new MqttChannelInitializer(connectFuture, host, port, mqttClientConfig.getSslContext()));
        ChannelFuture connectFut = bootstrap.connect();
        connectFut.addListener((ChannelFutureListener) f -> {
            if (f.isSuccess()) {
                this.channel = f.channel();
                log.debug("[{}][{}] Connected successfully {}!", host, port, this.channel.id());
                this.channel.closeFuture().addListener((ChannelFutureListener) cf -> {
                    if (isConnected()) {
                        return;
                    }
                    log.debug("[{}][{}] Channel is closed {}!", host, port, this.channel.id());
                    if (this.callback != null) {
                        this.callback.connectionLost(new RuntimeException("channel is closed!"));
                    }
                    //todo
                    pendingSubscriptions.forEach((id, mqttPendingSubscription) -> {
                        mqttPendingSubscription.onChannelClosed();
                    });
                    pendingSubscriptions.clear();
                    serverSubscriptions.clear();
                    subscriptions.clear();
                    pendingServerUnsubscribes.forEach((id, mqttPendingServerUnsubscribes) -> {
                        mqttPendingServerUnsubscribes.onChannelClosed();
                    });
                    pendingServerUnsubscribes.clear();
                    qos2PendingIncomingPublishes.clear();
                    pendingPublishes.forEach((id, mqttPendingPublish) -> {
                        mqttPendingPublish.onChannelClosed();
                    });
                    pendingPublishes.clear();
                    pendingSubscribeTopics.clear();
                    handlerToSubscription.clear();
                    this.scheduleConnectIfRequired(host, port, true);
                });
            } else {
                log.debug("[{}][{}] Connect failed, trying reconnect!", host, port);
                this.scheduleConnectIfRequired(host, port, reconnect);
            }
        });
        this.host = host;
        this.port = port;
        return connectFuture;
    }

    private void scheduleConnectIfRequired(String host, int port, boolean reconnect) {

        log.trace("[{}] Scheduling connect to server, isReconnect - {}", channel != null ? channel.id() : "UNKNOWN", reconnect);
        if (mqttClientConfig.isReconnect() && !disconnected) {
            if (reconnect) {
                this.reconnect = true;
            }
            eventLoop.schedule((Runnable) () -> connect(host, port, reconnect), mqttClientConfig.getReconnectDelay(), TimeUnit.SECONDS);
        }
    }

    @Override
    public boolean isConnected() {
        return channel != null && channel.isActive() && !this.disconnected;
    }

    @Override
    public Future<MqttConnectResult> reconnect() {

        log.trace("[{}] Reconnecting to server, isReconnect - {}", channel != null ? channel.id() : "UNKNOWN", reconnect);
        if (host == null) {
            throw new IllegalStateException("Cannot reconnect. host is null");
        }
        return connect(host, port);
    }

    /**
     * By default we use the netty {@link NioEventLoopGroup}.
     *
     * @return
     */
    @Override
    public EventLoopGroup getEventLoop() {
        return this.eventLoop;
    }

    @Override
    public void setEventLoop(EventLoopGroup eventLoop) {
        this.eventLoop = eventLoop;
    }

    @Override
    public Future<Void> on(String topic, MqttHandler handler) {
        return on(topic, handler, MqttQoS.AT_MOST_ONCE);
    }

    /**
     * 订阅指定的主题, 和QOS策略, 当接收到消息时, MqttClient将调用该消息处理器{@link MqttHandler}
     *
     * @param topic   订阅的消息 topic
     * @param handler 接收到消息时的消息处理器
     * @param qos     消息请求 qos
     * @return
     */
    @Override
    public Future<Void> on(String topic, MqttHandler handler, MqttQoS qos) {
        return createSubscription(topic, handler, false, qos);
    }

    @Override
    public Future<Void> once(String topic, MqttHandler handler) {
        return once(topic, handler, MqttQoS.AT_MOST_ONCE);
    }

    /**
     * 订阅指定的主题, 和QOS策略, 当接收到消息时, MqttClient将调用该消息处理器{@link MqttHandler}
     * 此订阅仅一次, 如果MqttClient收到了1条消息, 那么订阅将被删除
     *
     * @param topic   订阅的消息 topic
     * @param handler 接收到消息时的消息处理器
     * @param qos     消息请求 qos
     * @return
     */
    @Override
    public Future<Void> once(String topic, MqttHandler handler, MqttQoS qos) {
        return createSubscription(topic, handler, true, qos);
    }

    /**
     * 移除指定topic 或 handler 的订阅
     *
     * @param topic
     * @param handler
     * @return
     */
    @Override
    public Future<Void> off(String topic, MqttHandler handler) {

        log.trace("[{}] Unsubscribing from {}", channel != null ? channel.id() : "UNKNOWN", topic);
        Promise<Void> future = new DefaultPromise<>(this.eventLoop.next());
        for (MqttSubscription subscription : this.handlerToSubscription.get(handler)) {
            this.subscriptions.remove(topic, subscription);
        }
        this.handlerToSubscription.removeAll(handler);
        this.checkSubscriptions(topic, future);
        return future;
    }

    /**
     * 移除对应 topic 的所有订阅
     *
     * @param topic
     * @return
     */
    @Override
    public Future<Void> off(String topic) {

        log.trace("[{}] Unsubscribing from {}", channel != null ? channel.id() : "UNKNOWN", topic);
        Promise<Void> future = new DefaultPromise<>(this.eventLoop.next());
        ImmutableSet<MqttSubscription> subscriptions = ImmutableSet.copyOf(this.subscriptions.get(topic));
        for (MqttSubscription subscription : subscriptions) {
            for (MqttSubscription handSub : this.handlerToSubscription.get(subscription.getHandler())) {
                this.subscriptions.remove(topic, handSub);
            }
            this.handlerToSubscription.remove(subscription.getHandler(), subscription);
        }
        this.checkSubscriptions(topic, future);
        return future;
    }

    /**
     * 消息发布
     *
     * @param topic   主题
     * @param payload 内容
     * @return
     */
    @Override
    public Future<Void> publish(String topic, ByteBuf payload) {
        return this.publish(topic, payload, MqttQoS.AT_MOST_ONCE, false);
    }

    /**
     * 消息发布
     *
     * @param topic   主题
     * @param payload 内容
     * @param qos     qos
     * @return
     */
    @Override
    public Future<Void> publish(String topic, ByteBuf payload, MqttQoS qos) {
        return this.publish(topic, payload, qos, false);
    }

    @Override
    public Future<Void> publish(String topic, ByteBuf payload, boolean retain) {
        return this.publish(topic, payload, MqttQoS.AT_MOST_ONCE, retain);
    }

    /**
     * 发布消息
     *
     * @param topic   目标主题
     * @param payload 内容
     * @param qos     qos
     * @param retain  是否持久化
     * @return
     */
    @Override
    public Future<Void> publish(String topic, ByteBuf payload, MqttQoS qos, boolean retain) {

        log.trace("[{}] Publishing message to {}", channel != null ? channel.id() : "UNKNOWN", topic);
        Promise<Void> future = new DefaultPromise<>(this.eventLoop.next());
        //构造发布消息体
        MqttFixedHeader fixedHeader = new MqttFixedHeader(MqttMessageType.PUBLISH, false, qos, retain, 0);
        MqttPublishVariableHeader variableHeader = new MqttPublishVariableHeader(topic, getNewMessageId().messageId());
        MqttPublishMessage message = new MqttPublishMessage(fixedHeader, variableHeader, payload);

        MqttPendingPublish pendingPublish = new MqttPendingPublish(variableHeader.packetId(), future, message, qos,
                () -> pendingSubscriptions.containsKey(variableHeader.packetId()));
        this.pendingPublishes.put(pendingPublish.getMessageId(), pendingPublish);
        ChannelFuture channelFuture = this.sendAndFlushPacket(message);
        if (channelFuture == null) {
            pendingPublishes.remove(pendingPublish.getMessageId());
            return future;
        }
        channelFuture.addListener(result -> {
            //已寄送
            pendingPublish.setSent(true);
            if (result.cause() != null) {
                pendingPublishes.remove(pendingPublish.getMessageId());
                future.setFailure(result.cause());
            } else {
                if (pendingPublish.isSent() && pendingPublish.getQos() == MqttQoS.AT_MOST_ONCE) {
                    //ack for qos 0
                    pendingPublishes.remove(pendingPublish.getMessageId());
                    pendingPublish.getFuture().setSuccess(null);
                } else if (pendingPublish.isSent()) {
                    pendingPublish.startPublishRetransmissionTimer(eventLoop.next(), this::sendAndFlushPacket);
                } else {
                    pendingPublishes.remove(pendingPublish.getMessageId());
                }
            }
        });
        return future;
    }

    @Override
    public MqttClientConfig getClientConfig() {
        return this.mqttClientConfig;
    }

    @Override
    public void disconnect() {

        log.trace("[{}] Disconnecting from server", channel != null ? channel.id() : "UNKNOWN");
        disconnected = true;
        if (this.channel != null) {
            MqttFixedHeader fixedHeader = new MqttFixedHeader(MqttMessageType.DISCONNECT, false, MqttQoS.AT_MOST_ONCE, false, 0);
            MqttMessage message = new MqttMessage(fixedHeader);
            this.sendAndFlushPacket(message)
                    .addListener(future1 -> {
                        channel.close();
                    });
        }
    }

    @Override
    public void setCallback(MqttClientCallback callback) {
        this.callback = callback;
    }

    public boolean isReconnect() {
        return this.reconnect;
    }

    private Future<Void> createSubscription(String topic, MqttHandler handler, boolean once, MqttQoS qos) {

        log.trace("[{}] Creating subscription to {}", channel != null ? channel.id() : "UNKNOWN", topic);
        if (this.pendingSubscribeTopics.contains(topic)) {
            Optional<Map.Entry<Integer, MqttPendingSubscription>> subscriptionEntry = this.pendingSubscriptions.entrySet()
                    .stream()
                    .filter(e -> e.getValue().getTopic().equals(topic))
                    .findAny();
            if (subscriptionEntry.isPresent()) {
                MqttPendingSubscription pendingSubscription = subscriptionEntry.get().getValue();
                pendingSubscription.addHandler(handler, once);
                return pendingSubscription.getFuture();
            }
        }
        if (this.serverSubscriptions.contains(topic)) {

            MqttSubscription subscription = new MqttSubscription(topic, handler, once);
            this.subscriptions.put(topic, subscription);
            this.handlerToSubscription.put(handler, subscription);
            return this.channel.newSucceededFuture();
        }
        Promise<Void> future = new DefaultPromise<>(this.eventLoop.next());
        MqttFixedHeader fixedHeader = new MqttFixedHeader(MqttMessageType.SUBSCRIBE, false, MqttQoS.AT_LEAST_ONCE, false, 0);
        MqttTopicSubscription subscription = new MqttTopicSubscription(topic, qos);
        MqttMessageIdVariableHeader variableHeader = getNewMessageId();
        MqttSubscribePayload payload = new MqttSubscribePayload(Collections.singletonList(subscription));
        MqttSubscribeMessage message = new MqttSubscribeMessage(fixedHeader, variableHeader, payload);

        final MqttPendingSubscription pendingSubscription = new MqttPendingSubscription(future, topic, message,
                () -> !pendingSubscriptions.containsKey(variableHeader.messageId()));
        pendingSubscription.addHandler(handler, once);
        this.pendingSubscriptions.put(variableHeader.messageId(), pendingSubscription);
        this.pendingSubscribeTopics.add(topic);
        pendingSubscription.setSent(sendAndFlushPacket(message) != null);
        pendingSubscription.startRetransmitTimer(eventLoop.next(), this::sendAndFlushPacket);
        return future;
    }

    private void checkSubscriptions(String topic, Promise<Void> promise) {
        if (!(this.subscriptions.containsKey(topic) && this.subscriptions.get(topic).size() != 0) && this.serverSubscriptions.contains(topic)) {
            MqttFixedHeader fixedHeader = new MqttFixedHeader(MqttMessageType.UNSUBSCRIBE, false, MqttQoS.AT_LEAST_ONCE, false, 0);
            MqttMessageIdVariableHeader variableHeader = getNewMessageId();
            MqttUnsubscribePayload payload = new MqttUnsubscribePayload(Collections.singletonList(topic));
            MqttUnsubscribeMessage message = new MqttUnsubscribeMessage(fixedHeader, variableHeader, payload);

            MqttPendingUnsubscription pendingUnsubscription = new MqttPendingUnsubscription(promise, topic, message,
                    () -> !pendingServerUnsubscribes.containsKey(variableHeader.messageId()));
            this.pendingServerUnsubscribes.put(variableHeader.messageId(), pendingUnsubscription);
            pendingUnsubscription.startRetransmissionTimer(this.eventLoop.next(), this::sendAndFlushPacket);

            this.sendAndFlushPacket(message);
        } else {
            promise.setSuccess(null);
        }
    }

    ChannelFuture sendAndFlushPacket(Object message) {
        if (this.channel == null) {
            return null;
        }
        if (this.channel.isActive()) {
            log.trace("[{}] Sending message {}", channel != null ? channel.id() : "UNKNOWN", message);
            return this.channel.writeAndFlush(message);
        }
        return this.channel.newFailedFuture(new RuntimeException("Channel is closed!"));
    }

    private MqttMessageIdVariableHeader getNewMessageId() {
        int messageId;
        synchronized (this.nextMessageId) {
            this.nextMessageId.compareAndSet(0xffff, 1);
            messageId = this.nextMessageId.getAndIncrement();
        }
        return MqttMessageIdVariableHeader.from(messageId);
    }

    ConcurrentMap<Integer, MqttIncomingQos2Publish> getQos2PendingIncomingPublishes() {
        return qos2PendingIncomingPublishes;
    }

    private class MqttChannelInitializer extends ChannelInitializer<SocketChannel> {

        private final Promise<MqttConnectResult> connectFuture;
        private final String host;
        private final int port;
        private final SslContext sslContext;


        public MqttChannelInitializer(Promise<MqttConnectResult> connectFuture, String host, int port, SslContext sslContext) {
            this.connectFuture = connectFuture;
            this.host = host;
            this.port = port;
            this.sslContext = sslContext;
        }

        @Override
        protected void initChannel(SocketChannel ch) throws Exception {
            if (sslContext != null) {
                ch.pipeline().addLast(sslContext.newHandler(ch.alloc(), host, port));
            }

            ch.pipeline().addLast("mqttDecoder", new MqttDecoder(mqttClientConfig.getMaxBytesInMessage()));
            ch.pipeline().addLast("mqttEncoder", MqttEncoder.INSTANCE);
            ch.pipeline().addLast("idleStateHandler", new IdleStateHandler(mqttClientConfig.getTimeoutSeconds(), mqttClientConfig.getTimeoutSeconds(), 0));
            ch.pipeline().addLast("mqttPingHandler", new MqttPingHandler(mqttClientConfig.getTimeoutSeconds()));
            ch.pipeline().addLast("mqttHandler", new MqttChannelHandler(MqttClientImpl.this, connectFuture));
        }
    }
}
